// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/common/parallel_for.h>
#include <pollux/common/base/exceptions.h>
#include <pollux/dwio/common/executor_barrier.h>

namespace kumo::pollux::dwio::common {
    namespace {
        std::vector<std::pair<size_t, size_t> >
        splitRange(size_t from, size_t to, size_t factor) {
            POLLUX_CHECK_LE(from, to);
            std::vector<std::pair<size_t, size_t> > ranges;

            if (from == to) {
                return ranges;
            }

            if (factor <= 1) {
                ranges.emplace_back(from, to);
                return ranges;
            }

            auto rangeSize = to - from;
            auto chunkSize = rangeSize / factor;
            auto remainder = rangeSize % factor;
            auto start = from;
            for (size_t i = 0; i < factor; ++i) {
                auto end = start + chunkSize;
                if (remainder > 0) {
                    --remainder;
                    ++end;
                }
                // If `factor > (to - from)`, the rest of the chunks will be empty
                if (end > start) {
                    ranges.emplace_back(start, end);
                } else {
                    break;
                }
                start = end;
            }
            return ranges;
        }
    } // namespace

    ParallelFor::ParallelFor(
        melon::Executor::KeepAlive<> executor,
        size_t from,
        size_t to,
        size_t parallelismFactor)
        : executor_(executor),
          ranges_{splitRange(from, to, (executor_ ? parallelismFactor : 0))} {
    }

    ParallelFor::ParallelFor(
        std::shared_ptr<melon::Executor> executor,
        size_t from,
        size_t to,
        size_t parallelismFactor)
        : ParallelFor{executor.get(), from, to, parallelismFactor} {
        owned_ = std::move(executor);
    }

    void ParallelFor::execute(std::function<void(size_t)> func) {
        // Otherwise from == to
        if (ranges_.empty()) {
            return;
        }
        if (ranges_.size() == 1) {
            for (size_t i = ranges_[0].first, end = ranges_[0].second; i < end; ++i) {
                func(i);
            }
        } else {
            POLLUX_CHECK(
                executor_,
                "Executor wasn't provided so we shouldn't have more than 1 range");
            ExecutorBarrier barrier(executor_);
            const size_t last = ranges_.size() - 1;
            // First N-1 ranges in executor threads
            for (size_t r = 0; r < last; ++r) {
                auto &range = ranges_[r];
                barrier.add([begin = range.first, end = range.second, &func]() {
                    for (size_t i = begin; i < end; ++i) {
                        func(i);
                    }
                });
            }
            // Last range in calling thread
            auto &range = ranges_[last];
            for (size_t i = range.first, end = range.second; i < end; ++i) {
                func(i);
            }
            barrier.waitAll();
        }
    }
} // namespace kumo::pollux::dwio::common
